#!/usr/bin/env python3
"""
对话记录MCP服务器
提供对话记录的增删改查功能，确保只存储必要的对话信息
"""

import os
import sys
import json
import sqlite3
import uuid
import datetime
from typing import Dict, List, Any, Optional
import logging
from enum import Enum

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stderr),
        logging.FileHandler('conversation_mcp_server.log')
    ]
)
logger = logging.getLogger("conversation-mcp")

# 数据库路径
DB_PATH = os.environ.get('DB_PATH', os.path.join(os.path.dirname(os.path.abspath(__file__)), 'conversations.db'))

class MessageType(Enum):
    USER_MESSAGE = 'user_message'
    ASSISTANT_RESPONSE = 'assistant_response'
    SYSTEM_MESSAGE = 'system_message'
    TEMPORARY = 'temporary'

class ConversationDB:
    def __init__(self, db_path: str):
        """初始化数据库连接"""
        self.db_path = db_path
        self._init_db()

    def _init_db(self):
        """初始化数据库表结构"""
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS conversations (
                    id TEXT PRIMARY KEY,
                    datetime TEXT NOT NULL,
                    topic TEXT NOT NULL,
                    keywords TEXT NOT NULL,
                    summary TEXT NOT NULL,
                    message_type TEXT NOT NULL,
                    content TEXT NOT NULL,
                    metadata TEXT,
                    status TEXT NOT NULL DEFAULT 'active'
                )
            """)
            conn.execute("CREATE INDEX IF NOT EXISTS idx_conversations_datetime ON conversations(datetime)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_conversations_status ON conversations(status)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_conversations_message_type ON conversations(message_type)")

    def _validate_message(self, content: Dict[str, Any]) -> bool:
        """验证消息内容是否符合要求"""
        required_fields = ['message_type', 'content']
        if not all(field in content for field in required_fields):
            return False
        
        try:
            message_type = MessageType(content['message_type'])
            if message_type == MessageType.TEMPORARY:
                return False
        except ValueError:
            return False
        
        return True

    def add_conversation(self, content: Dict[str, Any]) -> str:
        """添加新的对话记录"""
        if not self._validate_message(content):
            raise ValueError("Invalid message content")

        conversation_id = str(uuid.uuid4())
        now = datetime.datetime.now().isoformat()

        # 过滤掉临时信息和系统提示
        if content.get('message_type') in [MessageType.TEMPORARY.value, MessageType.SYSTEM_MESSAGE.value]:
            logger.info(f"Skipping temporary or system message: {content.get('message_type')}")
            return None

        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                INSERT INTO conversations (
                    id, datetime, topic, keywords, summary, 
                    message_type, content, metadata, status
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
            """, (
                conversation_id,
                now,
                content.get('topic', ''),
                json.dumps(content.get('keywords', [])),
                content.get('summary', ''),
                content['message_type'],
                content['content'],
                json.dumps(content.get('metadata', {})),
                'active'
            ))
        return conversation_id

    def get_conversation(self, conversation_id: str) -> Optional[Dict[str, Any]]:
        """获取特定对话记录"""
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            cursor = conn.execute(
                "SELECT * FROM conversations WHERE id = ? AND status = 'active'",
                (conversation_id,)
            )
            row = cursor.fetchone()
            if row:
                return dict(row)
        return None

    def list_conversations(self, limit: int = 100, offset: int = 0) -> List[Dict[str, Any]]:
        """列出对话记录"""
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            cursor = conn.execute(
                """
                SELECT * FROM conversations 
                WHERE status = 'active' 
                AND message_type IN (?, ?)
                ORDER BY datetime DESC LIMIT ? OFFSET ?
                """,
                (MessageType.USER_MESSAGE.value, MessageType.ASSISTANT_RESPONSE.value, limit, offset)
            )
            return [dict(row) for row in cursor.fetchall()]

    def delete_conversation(self, conversation_id: str) -> bool:
        """删除对话记录（软删除）"""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.execute(
                "UPDATE conversations SET status = 'deleted' WHERE id = ?",
                (conversation_id,)
            )
            return cursor.rowcount > 0

    def cleanup_temporary_data(self) -> int:
        """清理临时数据"""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.execute(
                "DELETE FROM conversations WHERE message_type = ?",
                (MessageType.TEMPORARY.value,)
            )
            return cursor.rowcount

class ConversationMCPServer:
    def __init__(self):
        self.db = ConversationDB(DB_PATH)
        self.stdin = os.fdopen(sys.stdin.fileno(), 'rb', 0)
        self.stdout = os.fdopen(sys.stdout.fileno(), 'wb', 0)

    def read_message(self) -> Optional[Dict[str, Any]]:
        """从标准输入读取消息"""
        try:
            header = self.stdin.read(4)
            if not header or len(header) != 4:
                return None
            
            length = int.from_bytes(header, byteorder='big')
            message_bytes = self.stdin.read(length)
            if not message_bytes or len(message_bytes) != length:
                return None
            
            return json.loads(message_bytes.decode('utf-8'))
        except Exception as e:
            logger.error(f"读取消息失败: {e}")
            return None

    def write_message(self, message: Dict[str, Any]) -> None:
        """向标准输出写入消息"""
        try:
            message_bytes = json.dumps(message).encode('utf-8')
            length = len(message_bytes)
            header = length.to_bytes(4, byteorder='big')
            self.stdout.write(header)
            self.stdout.write(message_bytes)
            self.stdout.flush()
        except Exception as e:
            logger.error(f"写入消息失败: {e}")
            raise

    def handle_request(self, request: Dict[str, Any]) -> Dict[str, Any]:
        """处理请求"""
        method = request.get('method')
        params = request.get('params', {})

        if method == 'add_conversation':
            conversation_id = self.db.add_conversation(params)
            return {'result': {'id': conversation_id}}
        
        elif method == 'get_conversation':
            conversation = self.db.get_conversation(params.get('id'))
            return {'result': conversation}
        
        elif method == 'list_conversations':
            conversations = self.db.list_conversations(
                params.get('limit', 100),
                params.get('offset', 0)
            )
            return {'result': conversations}
        
        elif method == 'delete_conversation':
            success = self.db.delete_conversation(params.get('id'))
            return {'result': {'success': success}}
        
        elif method == 'cleanup':
            count = self.db.cleanup_temporary_data()
            return {'result': {'cleaned_count': count}}
        
        else:
            return {'error': {'code': 'method_not_found', 'message': f'Unknown method: {method}'}}

    def run(self):
        """运行服务器"""
        logger.info("对话记录MCP服务器启动")
        while True:
            try:
                message = self.read_message()
                if message is None:
                    continue

                response = self.handle_request(message)
                self.write_message(response)
            except Exception as e:
                logger.error(f"处理请求失败: {e}")
                error_response = {
                    'error': {
                        'code': 'internal_error',
                        'message': str(e)
                    }
                }
                try:
                    self.write_message(error_response)
                except:
                    pass

def main():
    try:
        server = ConversationMCPServer()
        server.run()
    except Exception as e:
        logger.error(f"服务器启动失败: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()